import time
import json
import random
import re
import websocket
from urllib.parse import urlparse
from locust import HttpUser, task, between, events

class WebSocketClient:
    def __init__(self, host):
        self.host = host
        self.ws = None
        self.connected = False
        self.user_id = None
        self.token = None
        self.character_id = None
        self.ws_url = None
    
    def extract_websocket_url(self, html_content):
        """从HTML内容中提取WebSocket连接URL"""
        # 这里需要根据实际游戏的HTML结构来提取WebSocket URL
        # 以下是一个示例，实际情况可能需要调整
        ws_pattern = r'(wss?://[^"\']+)'
        match = re.search(ws_pattern, html_content)
        if match:
            return match.group(1)
        return None
    
    def connect(self, user_id=None, token=None, html_content=None):
        """建立WebSocket连接"""
        try:
            # 如果提供了HTML内容，尝试从中提取WebSocket URL
            if html_content:
                self.ws_url = self.extract_websocket_url(html_content)
                print(f"从HTML中提取的WebSocket URL: {self.ws_url}")
            
            # 如果没有提取到WebSocket URL，使用默认构造
            if not self.ws_url:
                parsed_url = urlparse(self.host)
                ws_protocol = "wss" if parsed_url.scheme == "https" else "ws"
                ws_host = parsed_url.netloc
                self.ws_url = f"{ws_protocol}://{ws_host}/websocket"
            
            # 如果有登录凭证，则在URL中添加
            if user_id and token:
                if "?" in self.ws_url:
                    self.ws_url += f"&user_id={user_id}&token={token}"
                else:
                    self.ws_url += f"?user_id={user_id}&token={token}"
            
            print(f"尝试连接WebSocket: {self.ws_url}")
            self.ws = websocket.create_connection(self.ws_url)
            self.connected = True
            self.user_id = user_id
            self.token = token
            return True
        except Exception as e:
            print(f"WebSocket连接失败: {e}")
            return False
    
    def disconnect(self):
        """关闭WebSocket连接"""
        if self.ws and self.connected:
            self.ws.close()
            self.connected = False
    
    def send_message(self, message_type, data):
        """发送WebSocket消息"""
        if not self.connected:
            print("WebSocket未连接，无法发送消息")
            return False
        
        message = {
            "type": message_type,
            "data": data,
            "timestamp": int(time.time() * 1000)
        }
        
        try:
            self.ws.send(json.dumps(message))
            return True
        except Exception as e:
            print(f"发送消息失败: {e}")
            return False
    
    def receive_message(self, timeout=5):
        """接收WebSocket消息"""
        if not self.connected:
            print("WebSocket未连接，无法接收消息")
            return None
        
        try:
            self.ws.settimeout(timeout)
            response = self.ws.recv()
            return json.loads(response)
        except Exception as e:
            print(f"接收消息失败: {e}")
            return None


class GameUser(HttpUser):
    wait_time = between(1, 3)  # 用户操作之间的等待时间
    host = "https://cdn.xy.siyuf.com"  # 固定游戏服务器地址
    
    def on_start(self):
        """用户会话开始时执行"""
        self.ws_client = WebSocketClient(self.host)
        self.username = f"test_user_{random.randint(1000, 9999)}"
        self.password = "password123"
        self.logged_in = False
        self.in_game = False
        self.html_content = None
        
        # 首先访问游戏入口页面
        self.visit_game_page()
        
        # 注册新用户（如果需要）
        self.register()
        
        # 登录
        self.login()
        
        # 连接WebSocket
        if self.logged_in:
            self.connect_websocket()
            
            # 创建角色（如果需要）
            if not self.character_id:
                self.create_character()
            
            # 选择角色进入游戏
            if self.character_id:
                self.enter_game()
    
    def on_stop(self):
        """用户会话结束时执行"""
        # 登出游戏
        if self.in_game:
            self.logout_game()
        
        # 断开WebSocket连接
        if self.ws_client.connected:
            self.ws_client.disconnect()
    
    def visit_game_page(self):
        """访问游戏入口页面"""
        try:
            with self.client.get("/XYres/index_waiwang.html?isPCWeiduan=true", catch_response=True) as response:
                if response.status_code == 200:
                    self.html_content = response.text
                    print("成功访问游戏入口页面")
                    
                    # 提取可能需要的信息，如CSRF令牌等
                    # 这里需要根据实际游戏页面结构进行调整
                else:
                    print(f"访问游戏入口页面失败: HTTP {response.status_code}")
        except Exception as e:
            print(f"访问游戏入口页面时发生错误: {e}")
    
    def register(self):
        """注册新用户"""
        # 注意：这里需要根据实际游戏的注册API进行调整
        try:
            # 从HTML内容中提取注册所需的参数（如果有）
            register_data = {
                "username": self.username,
                "password": self.password,
                "email": f"{self.username}@example.com"
            }
            
            # 可能需要添加额外的参数，如验证码、邀请码等
            
            with self.client.post("/api/user/register", json=register_data, catch_response=True) as response:
                if response.status_code == 200:
                    try:
                        result = response.json()
                        if result.get("code") == 0 or result.get("success"):
                            print(f"用户 {self.username} 注册成功")
                        else:
                            print(f"用户注册失败: {result.get('message', '未知错误')}")
                    except Exception:
                        # 如果响应不是JSON格式，可能是其他格式
                        print(f"用户注册响应解析失败，响应内容: {response.text[:100]}...")
                else:
                    print(f"用户注册请求失败: HTTP {response.status_code}")
        except Exception as e:
            print(f"注册过程中发生错误: {e}")
            # 对于某些游戏，可能不需要注册，直接使用游客模式
            print("尝试使用游客模式继续...")
            self.logged_in = True  # 假设游客模式可以直接进入
    
    def login(self):
        """用户登录"""
        # 注意：这里需要根据实际游戏的登录API进行调整
        try:
            login_data = {
                "username": self.username,
                "password": self.password
            }
            
            # 可能需要添加额外的参数，如验证码等
            
            with self.client.post("/api/user/login", json=login_data, catch_response=True) as response:
                if response.status_code == 200:
                    try:
                        result = response.json()
                        if result.get("code") == 0 or result.get("success"):
                            # 提取用户ID和令牌，字段名可能需要调整
                            self.user_id = result.get("data", {}).get("userId") or result.get("user_id")
                            self.token = result.get("data", {}).get("token") or result.get("token")
                            self.logged_in = True
                            print(f"用户 {self.username} 登录成功")
                        else:
                            print(f"用户登录失败: {result.get('message', '未知错误')}")
                    except Exception:
                        # 如果响应不是JSON格式，可能是其他格式
                        print(f"用户登录响应解析失败，响应内容: {response.text[:100]}...")
                else:
                    print(f"用户登录请求失败: HTTP {response.status_code}")
        except Exception as e:
            print(f"登录过程中发生错误: {e}")
    
    def connect_websocket(self):
        """连接WebSocket"""
        if self.ws_client.connect(self.user_id, self.token, self.html_content):
            print(f"用户 {self.username} WebSocket连接成功")
        else:
            print(f"用户 {self.username} WebSocket连接失败")
    
    def create_character(self):
        """创建游戏角色"""
        character_name = f"hero_{random.randint(1000, 9999)}"
        character_class = random.choice(["warrior", "mage", "archer", "priest"])
        
        success = self.ws_client.send_message("create_character", {
            "name": character_name,
            "class": character_class,
            "appearance": {
                "hair": random.randint(1, 10),
                "face": random.randint(1, 10),
                "color": random.randint(1, 10)
            }
        })
        
        if success:
            # 等待服务器响应
            response = self.ws_client.receive_message()
            if response and response.get("type") == "character_created" and response.get("success"):
                self.character_id = response.get("data", {}).get("character_id")
                print(f"角色 {character_name} 创建成功，ID: {self.character_id}")
            else:
                print(f"角色创建失败: {response.get('message', '未知错误') if response else '无响应'}")
    
    def enter_game(self):
        """选择角色进入游戏"""
        if not self.character_id:
            print("没有可用的角色，无法进入游戏")
            return
        
        success = self.ws_client.send_message("enter_game", {
            "character_id": self.character_id
        })
        
        if success:
            # 等待服务器响应
            response = self.ws_client.receive_message()
            if response and response.get("type") == "game_entered" and response.get("success"):
                self.in_game = True
                print(f"角色 {self.character_id} 成功进入游戏")
            else:
                print(f"进入游戏失败: {response.get('message', '未知错误') if response else '无响应'}")
    
    def logout_game(self):
        """退出游戏"""
        success = self.ws_client.send_message("logout", {})
        if success:
            self.in_game = False
            print(f"角色 {self.character_id} 成功退出游戏")
    
    @task(3)
    def move_character(self):
        """移动角色"""
        if not self.in_game:
            return
        
        # 随机移动方向和距离
        direction = random.choice(["north", "south", "east", "west"])
        distance = random.randint(1, 10)
        
        success = self.ws_client.send_message("move", {
            "direction": direction,
            "distance": distance
        })
        
        if success:
            response = self.ws_client.receive_message(timeout=2)
            if response and response.get("type") == "move_result" and response.get("success"):
                new_position = response.get("data", {}).get("position", {})
                print(f"角色移动成功，新位置: X={new_position.get('x')}, Y={new_position.get('y')}")
            else:
                if response:
                    print(f"角色移动失败: {response.get('message', '未知错误')}")
    
    @task(1)
    def chat_message(self):
        """发送聊天消息"""
        if not self.in_game:
            return
        
        messages = [
            "你好！", "有人组队吗？", "这个游戏真好玩！", 
            "请问新手任务在哪里？", "有人知道怎么升级快吗？"
        ]
        
        success = self.ws_client.send_message("chat", {
            "channel": "world",
            "message": random.choice(messages)
        })
        
        if success:
            print(f"聊天消息发送成功")
    
    @task(2)
    def check_surroundings(self):
        """查看周围环境"""
        if not self.in_game:
            return
        
        success = self.ws_client.send_message("get_surroundings", {
            "range": random.randint(5, 20)
        })
        
        if success:
            response = self.ws_client.receive_message(timeout=2)
            if response and response.get("type") == "surroundings" and response.get("success"):
                entities = response.get("data", {}).get("entities", [])
                print(f"周围有 {len(entities)} 个实体")
            else:
                if response:
                    print(f"获取周围环境失败: {response.get('message', '未知错误')}")


# 自定义事件处理
@events.init.add_listener
def on_locust_init(environment, **kwargs):
    print("Locust初始化完成，准备开始测试...")

@events.test_start.add_listener
def on_test_start(environment, **kwargs):
    print("测试开始...")

@events.test_stop.add_listener
def on_test_stop(environment, **kwargs):
    print("测试结束...")


if __name__ == "__main__":
    # 本地测试用
    import sys
    if len(sys.argv) > 1 and sys.argv[1] == "debug":
        # 创建一个测试用户并执行完整流程
        user = GameUser(None)
        # 游戏服务器地址已在类中固定设置
        user.on_start()
        
        # 执行几次任务
        for _ in range(5):
            user.move_character()
            time.sleep(1)
            user.chat_message()
            time.sleep(1)
            user.check_surroundings()
            time.sleep(1)
        
        user.on_stop()